#include "echo.h"

#include <errno.h>
#include <sys/time.h>
#include <sys/wait.h>

// ================================================================================

#define BUF_LEN 1024

static void usage() {
  printf("Usage: echo-client <Echo-Server IP address> <Message> <Number of clients>\n");
}


static int write_all(int sock, char *buf, int len, sockaddr_in_t *dst) {
  while (len) {
    int res = sendto(sock, buf, len, 0, (sockaddr_t*)dst, sizeof(sockaddr_in_t));
    if (res < 0) {
      return res;
    }

    len -= res;
    buf += res;
  }
}


static int read_all(int sock, char *buf, int len, sockaddr_in_t *src) {
  while (len) {
    socklen_t addrlen = sizeof(sockaddr_in_t);
    
    int res = recvfrom(sock, buf, len, 0, (sockaddr_t*)src, &addrlen);
    if (res < 0) {
      return res;
    }

    len -= res;
    buf += res;
  }
}

int main(int argc, char *argv[]) {
  int sock;
  sockaddr_in_t svr_addr;
  char buf[BUF_LEN];
  char recv_buf[BUF_LEN];
  int msg_len;
  struct timeval timeout;
  int cli_num, i;
  int leader;
  pid_t *children;

  if (argc < 3) {
    usage();
    return 1;
  }

  if (inet_aton(argv[1], &svr_addr.sin_addr) == -1) {
    usage();
    return 1;
  }

  if (argc >= 4) {
    cli_num = atoi(argv[3]);
  } else {
    cli_num = 1;
  }

  children = 0;
  leader = 1;

  for (i = 0; i < cli_num; ++i) {
    pid_t child = fork();

    if (child < 0) {
      perror("EchoClient");
      return 1;
    }

    if (child == 0) {
      leader = 0;
      break;
    } else {
      if (!children) {
        children = malloc(cli_num*sizeof(pid_t));
        children[0] = child;
      } else {
        children[i] = child;
      }
    }
  }

  sock = socket(AF_INET, SOCK_DGRAM, 0);
  if (sock == -1) {
    perror("Failed to create a network socket");
    return 1;
  }

  svr_addr.sin_family = AF_INET;
  svr_addr.sin_port = htons(ECHO_PORT);

  timeout.tv_sec = 5;
  timeout.tv_usec = 0;
  setsockopt(sock, SOL_SOCKET, SO_RCVTIMEO, &timeout, sizeof(timeout));

  msg_len = snprintf(buf, BUF_LEN, "Client-%d: %s", i, argv[2]);

  if (write_all(sock, buf, msg_len, &svr_addr) < 0) {
    printf("Client-%d: Error while transmitting the message\n", i, strerror(errno));
    close(sock);
    return 1;
  }

  memset(buf, BUF_LEN, 0);

  if (read_all(sock, buf, msg_len, &svr_addr) < 0) {
    printf("Client-%d: error while reading a UDP datagram: %s\n", i, strerror(errno));
    close(sock);
    return 1;
  } else {
    printf("Client-%d receives: %s\n", i, buf);
  }  

  close(sock);

  if (leader) {
    for (i = 0; i < cli_num; ++i) {
      int s;
      waitpid(children[i], &s, 0);
    }

    free(children);
  }

  return 0;
}
